{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 翻译 Relay 代码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "from torch import fx\n",
    "from torch.nn import Module\n",
    "\n",
    "import tvm.testing\n",
    "from tvm.relax.frontend.torch import from_fx\n",
    "from tvm.relay.frontend import from_pytorch\n",
    "from tvm import relay\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.contrib.msc.core.frontend import translate\n",
    "from tvm.contrib.msc.framework.tvm import codegen as tvm_codegen\n",
    "from tvm.contrib.msc.core import utils as msc_utils\n",
    "\n",
    "def _valid_target(target):\n",
    "    if not target:\n",
    "        return target\n",
    "    if target == \"ignore\":\n",
    "        return None\n",
    "    if target == \"cuda\" and not tvm.cuda().exist:\n",
    "        return None\n",
    "    if isinstance(target, str):\n",
    "        target = tvm.target.Target(target)\n",
    "    return target\n",
    "\n",
    "\n",
    "def _run_relax(relax_mod, target, datas):\n",
    "    relax_mod = tvm.relax.transform.LegalizeOps()(relax_mod)\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        relax_exec = tvm.relax.build(relax_mod, target)\n",
    "        runnable = tvm.relax.VirtualMachine(relax_exec, tvm.cpu())\n",
    "    res = runnable[\"main\"](*datas)\n",
    "    if isinstance(res, tvm.runtime.NDArray):\n",
    "        return [res.asnumpy()]\n",
    "    return [e.asnumpy() for e in res]\n",
    "\n",
    "\n",
    "def verify_model(torch_model, input_info, opt_config=None, codegen_config=None, build_target=None):\n",
    "    \"\"\"Compare relax with relay\"\"\"\n",
    "\n",
    "    graph_model = fx.symbolic_trace(torch_model)\n",
    "    with torch.no_grad():\n",
    "        expected = from_fx(graph_model, input_info)\n",
    "    expected = tvm.relax.transform.CanonicalizeBindings()(expected)\n",
    "\n",
    "    # graph from relay\n",
    "    datas = [np.random.rand(*i[0]).astype(i[1]) for i in input_info]\n",
    "    torch_datas = [torch.from_numpy(i) for i in datas]\n",
    "    with torch.no_grad():\n",
    "        scripted_model = torch.jit.trace(torch_model, tuple(torch_datas)).eval()  # type: ignore\n",
    "    shape_list = [(\"input\" + str(idx), i) for idx, i in enumerate(input_info)]\n",
    "    relay_mod, params = from_pytorch(scripted_model, shape_list)\n",
    "    graph, weights = translate.from_relay(relay_mod, params, opt_config=opt_config)\n",
    "    # to relax\n",
    "    codegen_config = codegen_config or {}\n",
    "    codegen_config.update({\"explicit_name\": False, \"from_relay\": True})\n",
    "    mod = tvm_codegen.to_relax(graph, weights, codegen_config)\n",
    "    if build_target:\n",
    "        build_target = _valid_target(build_target)\n",
    "        if not build_target:\n",
    "            return\n",
    "        tvm_datas = [tvm.nd.array(i) for i in datas]\n",
    "        expected_res = _run_relax(expected, build_target, tvm_datas)\n",
    "        if not graph.get_inputs():\n",
    "            tvm_datas = []\n",
    "        res = _run_relax(mod, build_target, tvm_datas)\n",
    "        for exp_r, new_r in zip(expected_res, res):\n",
    "            tvm.testing.assert_allclose(exp_r, new_r, atol=1e-5, rtol=1e-5)\n",
    "    else:\n",
    "        tvm.ir.assert_structural_equal(mod, expected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_conv1d():\n",
    "    \"\"\"test relay to relax for conv1d\"\"\"\n",
    "\n",
    "    class Conv1D1(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.conv = torch.nn.Conv1d(3, 6, 7, bias=True)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.conv(data)\n",
    "\n",
    "    class Conv1D2(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.conv = torch.nn.Conv1d(3, 6, 7, bias=False)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.conv(data)\n",
    "\n",
    "    input_info = [([1, 3, 10], \"float32\")]\n",
    "    verify_model(Conv1D1(), input_info)\n",
    "    verify_model(Conv1D2(), input_info)\n",
    "\n",
    "\n",
    "def test_conv2d():\n",
    "    \"\"\"test relay to relax for conv2d\"\"\"\n",
    "\n",
    "    class Conv2D1(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.conv = torch.nn.Conv2d(3, 6, 7, bias=True)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.conv(data)\n",
    "\n",
    "    class Conv2D2(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.conv = torch.nn.Conv2d(3, 6, 7, bias=False)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.conv(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Conv2D1(), input_info)\n",
    "    verify_model(Conv2D2(), input_info)\n",
    "\n",
    "\n",
    "def test_linear():\n",
    "    \"\"\"test relay to relax for linear\"\"\"\n",
    "\n",
    "    class Dense1(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.linear = torch.nn.Linear(10, 7, bias=True)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.linear(data)\n",
    "\n",
    "    class Dense2(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.linear = torch.nn.Linear(10, 7, bias=False)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.linear(data)\n",
    "\n",
    "    class MatMul1(Module):\n",
    "        def forward(self, x, y):\n",
    "            return torch.matmul(x, y)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Dense1(), input_info, build_target=\"llvm\")\n",
    "    verify_model(Dense2(), input_info, build_target=\"llvm\")\n",
    "    verify_model(MatMul1(), [([10, 10], \"float32\"), ([10, 10], \"float32\")], build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_bmm():\n",
    "    \"\"\"test relay to relax for bmm\"\"\"\n",
    "\n",
    "    class BMM(Module):\n",
    "        def forward(self, x, y):\n",
    "            return torch.bmm(x, y)\n",
    "\n",
    "    input_info = [((4, 128, 256), \"float32\"), ((4, 256, 512), \"float32\")]\n",
    "    verify_model(BMM(), input_info, opt_config={\"opt_level\": 3})\n",
    "\n",
    "\n",
    "def test_baddbmm():\n",
    "    \"\"\"test relay to relax for baddbmm\"\"\"\n",
    "\n",
    "    class BAddBMM1(Module):\n",
    "        def forward(self, c, x, y):\n",
    "            return torch.baddbmm(c, x, y)\n",
    "\n",
    "    class BAddBMM2(Module):\n",
    "        def forward(self, c, x, y):\n",
    "            return torch.baddbmm(c, x, y, alpha=2, beta=0)\n",
    "\n",
    "    input_info = [\n",
    "        ((4, 128, 512), \"float32\"),\n",
    "        ((4, 128, 256), \"float32\"),\n",
    "        ((4, 256, 512), \"float32\"),\n",
    "    ]\n",
    "    verify_model(BAddBMM1(), input_info, opt_config={\"opt_level\": 3}, build_target=\"llvm\")\n",
    "    verify_model(BAddBMM2(), input_info, opt_config={\"opt_level\": 3}, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_relu():\n",
    "    \"\"\"test relay to relax for relu\"\"\"\n",
    "\n",
    "    class ReLU(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.relu = torch.nn.ReLU()\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.relu(data)\n",
    "\n",
    "    class ReLU1(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.nn.functional.relu(data)\n",
    "\n",
    "    input_info = [([10, 10], \"float32\")]\n",
    "    verify_model(ReLU(), input_info)\n",
    "    verify_model(ReLU1(), input_info)\n",
    "\n",
    "\n",
    "def test_relu6():\n",
    "    \"\"\"test relay to relax for relu6\"\"\"\n",
    "\n",
    "    class ReLU6(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.relu6 = torch.nn.ReLU6()\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.relu6(data)\n",
    "\n",
    "    input_info = [([10, 10], \"float32\")]\n",
    "    verify_model(ReLU6(), input_info)\n",
    "\n",
    "\n",
    "def test_maxpool2d():\n",
    "    \"\"\"test relay to relax for maxpool2d\"\"\"\n",
    "\n",
    "    class MaxPool2d(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.pool = torch.nn.MaxPool2d(kernel_size=[1, 1])\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.pool(data)\n",
    "\n",
    "    class MaxPool2d2(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.pool = torch.nn.MaxPool2d(kernel_size=[2, 2], dilation=[2, 3])\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.pool(data)\n",
    "\n",
    "    class MaxPool2d3(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.pool = torch.nn.MaxPool2d(kernel_size=[4, 4], padding=2, stride=2)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.pool(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(MaxPool2d(), input_info)\n",
    "    verify_model(MaxPool2d2(), input_info)\n",
    "    verify_model(MaxPool2d3(), input_info)\n",
    "\n",
    "\n",
    "def test_avgpool2d():\n",
    "    \"\"\"test relay to relax for avgpool2d\"\"\"\n",
    "\n",
    "    class AvgPool2d(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.pool = torch.nn.AvgPool2d(kernel_size=[1, 1])\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.pool(data)\n",
    "\n",
    "    class AvgPool2d2(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.pool = torch.nn.AvgPool2d(kernel_size=[4, 4], stride=2, padding=2, ceil_mode=True)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.pool(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(AvgPool2d(), input_info)\n",
    "    verify_model(AvgPool2d2(), input_info)\n",
    "\n",
    "\n",
    "def test_adaptive_avgpool2d():\n",
    "    \"\"\"test relay to relax for adaptive_avgpool2d\"\"\"\n",
    "\n",
    "    class AdaptiveAvgPool2d0(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.pool = torch.nn.AdaptiveAvgPool2d([10, 10])\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.pool(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(AdaptiveAvgPool2d0(), input_info)\n",
    "\n",
    "\n",
    "def test_flatten():\n",
    "    \"\"\"test relay to relax for flatten\"\"\"\n",
    "\n",
    "    class Flatten(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.f = torch.nn.Flatten(2, -1)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.f(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Flatten(), input_info, opt_config={\"opt_level\": 3}, build_target=\"llvm\")\n",
    "    verify_model(\n",
    "        torch.nn.Flatten(2, -1), input_info, opt_config={\"opt_level\": 3}, build_target=\"llvm\"\n",
    "    )\n",
    "\n",
    "\n",
    "def test_batchnorm2d():\n",
    "    \"\"\"test relay to relax for batchnorm2d\"\"\"\n",
    "\n",
    "    class BatchNorm2d(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.batchnorm = torch.nn.BatchNorm2d(3)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.batchnorm(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(BatchNorm2d(), input_info, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_embedding():\n",
    "    \"\"\"test relay to relax for embedding\"\"\"\n",
    "\n",
    "    class Embedding(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.embedding = torch.nn.Embedding(10, 3)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.embedding(data)\n",
    "\n",
    "    verify_model(Embedding(), [([4], \"int64\")])\n",
    "    verify_model(Embedding(), [([4, 5], \"int64\")])\n",
    "\n",
    "\n",
    "def test_layernorm():\n",
    "    \"\"\"test relay to relax for layernorm\"\"\"\n",
    "\n",
    "    class LayerNorm(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.layernorm = torch.nn.LayerNorm(10)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.layernorm(data)\n",
    "\n",
    "    input_info = [([1, 10, 10], \"float32\")]\n",
    "    verify_model(LayerNorm(), input_info)\n",
    "\n",
    "\n",
    "def test_functional_layernorm():\n",
    "    \"\"\"test relay to relax for functional_layernorm\"\"\"\n",
    "\n",
    "    class LayerNorm(Module):\n",
    "        def __init__(self, shape):\n",
    "            super().__init__()\n",
    "            self.weight = torch.nn.Parameter(torch.ones(shape))\n",
    "            self.bias = torch.nn.Parameter(torch.zeros(shape))\n",
    "\n",
    "        def forward(self, data):\n",
    "            return torch.nn.functional.layer_norm(\n",
    "                data, self.weight.shape, self.weight, self.bias, 1e-5\n",
    "            )\n",
    "\n",
    "    input_info = [([1, 10, 10], \"float32\")]\n",
    "    verify_model(LayerNorm((10)), input_info)\n",
    "\n",
    "\n",
    "def test_cross_entropy():\n",
    "    \"\"\"test relay to relax for cross_entropy\"\"\"\n",
    "\n",
    "    class CrossEntropy1(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.loss = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "        def forward(self, logits, targets):\n",
    "            return self.loss(logits, targets)\n",
    "\n",
    "    class CrossEntropy2(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.weight = torch.nn.Parameter(torch.ones((2,)))\n",
    "            self.loss = torch.nn.CrossEntropyLoss(weight=self.weight)\n",
    "\n",
    "        def forward(self, logits, targets):\n",
    "            return self.loss(logits, targets)\n",
    "\n",
    "    class CrossEntropy3(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.loss = torch.nn.CrossEntropyLoss(ignore_index=1, reduction=\"sum\")\n",
    "\n",
    "        def forward(self, logits, targets):\n",
    "            return self.loss(logits, targets)\n",
    "\n",
    "    input_info = [([3, 2], \"float32\"), ([3], \"int64\")]\n",
    "    verify_model(CrossEntropy1(), input_info, opt_config={\"opt_level\": 3}, build_target=\"llvm\")\n",
    "    verify_model(CrossEntropy2(), input_info, opt_config={\"opt_level\": 3}, build_target=\"llvm\")\n",
    "    verify_model(CrossEntropy3(), input_info, opt_config={\"opt_level\": 3}, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_functional_cross_entropy():\n",
    "    \"\"\"test relay to relax for functional_cross_entropy\"\"\"\n",
    "\n",
    "    class CrossEntropy(Module):\n",
    "        def forward(self, logits, targets):\n",
    "            return torch.nn.functional.cross_entropy(logits, targets)\n",
    "\n",
    "    input_info = [([3, 10], \"float32\"), ([3], \"int64\")]\n",
    "    verify_model(CrossEntropy(), input_info, opt_config={\"opt_level\": 3}, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_silu():\n",
    "    \"\"\"test relay to relax for silu\"\"\"\n",
    "\n",
    "    class SiLU(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.silu = torch.nn.SiLU()\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.silu(data)\n",
    "\n",
    "    class SiLU2(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.nn.functional.silu(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(SiLU(), input_info, build_target=\"llvm\")\n",
    "    verify_model(SiLU2(), input_info, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_groupnorm():\n",
    "    \"\"\"test relay to relax for groupnorm\"\"\"\n",
    "\n",
    "    class GroupNorm(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.groupnorm = torch.nn.GroupNorm(3, 3)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.groupnorm(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(GroupNorm(), input_info)\n",
    "\n",
    "\n",
    "def test_softmax():\n",
    "    \"\"\"test relay to relax for softmax\"\"\"\n",
    "\n",
    "    class Softmax(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.softmax = torch.nn.Softmax(dim=1)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.softmax(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Softmax(), input_info)\n",
    "\n",
    "\n",
    "def test_binary():\n",
    "    \"\"\"test relay to relax for binary\"\"\"\n",
    "\n",
    "    input_info1 = [([1, 3, 10, 10], \"float32\"), ([1, 3, 10, 10], \"float32\")]\n",
    "    input_info2 = [([1, 3, 10, 10], \"float32\")]\n",
    "\n",
    "    # Add\n",
    "    class Add1(Module):\n",
    "        def forward(self, lhs, rhs):\n",
    "            return lhs + rhs\n",
    "\n",
    "    class Add2(Module):\n",
    "        def forward(self, lhs):\n",
    "            return lhs + 1.0\n",
    "\n",
    "    verify_model(Add1(), input_info1, opt_config={\"opt_level\": 3})\n",
    "    verify_model(Add2(), input_info2, opt_config={\"opt_level\": 3})\n",
    "\n",
    "    # Sub\n",
    "    class Sub1(Module):\n",
    "        def forward(self, lhs, rhs):\n",
    "            return lhs - rhs\n",
    "\n",
    "    class Sub2(Module):\n",
    "        def forward(self, lhs):\n",
    "            return lhs - 1.0\n",
    "\n",
    "    verify_model(Sub1(), input_info1, opt_config={\"opt_level\": 3})\n",
    "    verify_model(Sub2(), input_info2, opt_config={\"opt_level\": 3})\n",
    "\n",
    "    # Mul\n",
    "    class Mul1(Module):\n",
    "        def forward(self, lhs, rhs):\n",
    "            return lhs * rhs\n",
    "\n",
    "    class Mul2(Module):\n",
    "        def forward(self, lhs):\n",
    "            return lhs * 1.0\n",
    "\n",
    "    verify_model(Mul1(), input_info1, opt_config={\"opt_level\": 3})\n",
    "    verify_model(Mul2(), input_info2)\n",
    "\n",
    "    # True div\n",
    "    class TrueDiv1(Module):\n",
    "        def forward(self, lhs, rhs):\n",
    "            return lhs / rhs\n",
    "\n",
    "    class TrueDiv2(Module):\n",
    "        def forward(self, lhs):\n",
    "            return lhs / 1.0\n",
    "\n",
    "    verify_model(TrueDiv1(), input_info1, opt_config={\"opt_level\": 3})\n",
    "    verify_model(TrueDiv2(), input_info2)\n",
    "\n",
    "    # Floor div\n",
    "    class FloorDiv1(Module):\n",
    "        def forward(self, lhs, rhs):\n",
    "            return lhs // rhs\n",
    "\n",
    "    class FloorDiv2(Module):\n",
    "        def forward(self, lhs):\n",
    "            return lhs // 1.0\n",
    "\n",
    "    verify_model(FloorDiv1(), input_info1, opt_config={\"opt_level\": 3})\n",
    "    verify_model(FloorDiv2(), input_info2, opt_config={\"opt_level\": 3})\n",
    "\n",
    "    # Power\n",
    "    class Power1(Module):\n",
    "        def forward(self, lhs, rhs):\n",
    "            return lhs**rhs\n",
    "\n",
    "    class Power2(Module):\n",
    "        def forward(self, lhs):\n",
    "            return lhs**1.0\n",
    "\n",
    "    verify_model(Power1(), input_info1, opt_config={\"opt_level\": 3})\n",
    "    verify_model(Power2(), input_info2, opt_config={\"opt_level\": 3})\n",
    "\n",
    "    # LT\n",
    "    class LT1(Module):\n",
    "        def forward(self, lhs, rhs):\n",
    "            return lhs < rhs\n",
    "\n",
    "    class LT2(Module):\n",
    "        def forward(self, lhs):\n",
    "            return lhs < 1.0\n",
    "\n",
    "    verify_model(LT1(), input_info1, opt_config={\"opt_level\": 3})\n",
    "    verify_model(LT2(), input_info2, opt_config={\"opt_level\": 3})\n",
    "\n",
    "\n",
    "def test_squeeze():\n",
    "    \"\"\"test relay to relax for squeeze\"\"\"\n",
    "\n",
    "    class Squeeze1(Module):\n",
    "        def forward(self, data):\n",
    "            return data.squeeze(1)\n",
    "\n",
    "    class Squeeze2(Module):\n",
    "        def forward(self, data):\n",
    "            return data.squeeze()\n",
    "\n",
    "    input_info = [([3, 1, 4, 1], \"float32\")]\n",
    "    verify_model(Squeeze1(), input_info)\n",
    "    verify_model(Squeeze2(), input_info)\n",
    "\n",
    "\n",
    "def test_unsqueeze():\n",
    "    \"\"\"test relay to relax for unsqueeze\"\"\"\n",
    "\n",
    "    class Unsqueeze1(Module):\n",
    "        def forward(self, data):\n",
    "            return data.unsqueeze(1)\n",
    "\n",
    "    class Unsqueeze2(Module):\n",
    "        def forward(self, data):\n",
    "            return data.unsqueeze(-1)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Unsqueeze1(), input_info)\n",
    "    verify_model(Unsqueeze2(), input_info)\n",
    "\n",
    "\n",
    "def test_getitem():\n",
    "    \"\"\"test relay to relax for getitem\"\"\"\n",
    "\n",
    "    class Slice1(Module):\n",
    "        def forward(self, x):\n",
    "            return x[0, 1::2, :, :3]\n",
    "\n",
    "    class Slice2(Module):\n",
    "        def forward(self, x):\n",
    "            return x[:, None, None, :, None]\n",
    "\n",
    "    verify_model(Slice1(), [([1, 3, 10, 10], \"float32\")], build_target=\"ignore\")\n",
    "    verify_model(Slice2(), [([8, 16], \"float32\")], build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_unary():\n",
    "    \"\"\"test relay to relax for unary\"\"\"\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "\n",
    "    # sin\n",
    "    class Sin(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.sin(data)\n",
    "\n",
    "    verify_model(Sin(), input_info)\n",
    "\n",
    "    # cos\n",
    "    class Cos(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.cos(data)\n",
    "\n",
    "    verify_model(Cos(), input_info)\n",
    "\n",
    "    # exp\n",
    "    class Exp(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.exp(data)\n",
    "\n",
    "    verify_model(Exp(), input_info)\n",
    "\n",
    "    # sqrt\n",
    "    class Sqrt(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.sqrt(data)\n",
    "\n",
    "    verify_model(Sqrt(), input_info)\n",
    "\n",
    "    # sigmoid\n",
    "    class Sigmoid(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.sigmoid(data)\n",
    "\n",
    "    verify_model(Sigmoid(), input_info)\n",
    "\n",
    "    # round\n",
    "    class Round(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.round(data)\n",
    "\n",
    "    verify_model(Round(), input_info)\n",
    "\n",
    "\n",
    "def test_gelu():\n",
    "    \"\"\"test relay to relax for gelu\"\"\"\n",
    "\n",
    "    class Gelu(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.nn.functional.gelu(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Gelu(), input_info)\n",
    "\n",
    "\n",
    "def test_tanh():\n",
    "    \"\"\"test relay to relax for tanh\"\"\"\n",
    "\n",
    "    class Tanh(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.tanh(data)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Tanh(), input_info)\n",
    "\n",
    "\n",
    "def test_clamp():\n",
    "    \"\"\"test relay to relax for clamp\"\"\"\n",
    "\n",
    "    class Clamp(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.clamp(data, min=0.1, max=0.5)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Clamp(), input_info)\n",
    "\n",
    "\n",
    "def test_interpolate():\n",
    "    \"\"\"test relay to relax for interpolate\"\"\"\n",
    "\n",
    "    class Interpolate(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.nn.functional.interpolate(data, (5, 5))\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Interpolate(), input_info, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_addmm():\n",
    "    \"\"\"test relay to relax for addmm\"\"\"\n",
    "\n",
    "    class Addmm(Module):\n",
    "        def forward(self, x_1, x_2, x_3):\n",
    "            return torch.addmm(x_1, x_2, x_3)\n",
    "\n",
    "    input_info = [\n",
    "        ([10, 10], \"float32\"),\n",
    "        ([10, 10], \"float32\"),\n",
    "        ([10, 10], \"float32\"),\n",
    "    ]\n",
    "    verify_model(Addmm(), input_info, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_split():\n",
    "    \"\"\"test relay to relax for split\"\"\"\n",
    "\n",
    "    class Split(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.split(data, 1, dim=1)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Split(), input_info, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_cumsum():\n",
    "    \"\"\"test relay to relax for cumsum\"\"\"\n",
    "\n",
    "    class Cumsum(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.cumsum(data, dim=1, dtype=torch.int32)\n",
    "\n",
    "    input_info = [([1, 2, 3, 4], \"float32\")]\n",
    "    verify_model(Cumsum(), input_info)\n",
    "\n",
    "\n",
    "def test_chunk():\n",
    "    \"\"\"test relay to relax for chunk\"\"\"\n",
    "\n",
    "    class Chunk(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.chunk(data, 3, dim=1)\n",
    "\n",
    "    input_info = [([1, 3, 10, 10], \"float32\")]\n",
    "    verify_model(Chunk(), input_info, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_inplace_fill():\n",
    "    \"\"\"test relay to relax for inplace_fill\"\"\"\n",
    "\n",
    "    class InplaceFill(Module):\n",
    "        def forward(self, data):\n",
    "            data.fill_(1.5)\n",
    "            return data\n",
    "\n",
    "    verify_model(InplaceFill(), [([10, 10], \"float32\")], build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_arange():\n",
    "    \"\"\"test relay to relax for arange\"\"\"\n",
    "\n",
    "    class Arange(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.arange(0, 20, dtype=torch.int32)\n",
    "\n",
    "    verify_model(\n",
    "        Arange(), [([10, 10], \"float32\")], opt_config={\"opt_level\": 3}, build_target=\"llvm\"\n",
    "    )\n",
    "\n",
    "\n",
    "def test_empty():\n",
    "    \"\"\"test relay to relax for empty\"\"\"\n",
    "\n",
    "    class Empty(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.empty((10, 10), dtype=torch.float32)\n",
    "\n",
    "    verify_model(\n",
    "        Empty(), [([10, 10], \"float32\")], opt_config={\"opt_level\": 3}, build_target=\"ignore\"\n",
    "    )\n",
    "\n",
    "\n",
    "def test_tensor():\n",
    "    \"\"\"test relay to relax for tensor\"\"\"\n",
    "\n",
    "    class Empty1(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.tensor(3, dtype=torch.float32)\n",
    "\n",
    "    class Empty2(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.tensor(3)\n",
    "\n",
    "    verify_model(Empty1(), [([10, 10], \"float32\")], build_target=\"llvm\")\n",
    "    verify_model(Empty2(), [([10, 10], \"float32\")], build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_tril():\n",
    "    \"\"\"test relay to relax for tril\"\"\"\n",
    "\n",
    "    class Tril(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.tril(data, 1)\n",
    "\n",
    "    class InplaceTril(Module):\n",
    "        def forward(self, data):\n",
    "            data.tril_(1)\n",
    "            return data\n",
    "\n",
    "    input_info = [([10, 10], \"float32\")]\n",
    "    verify_model(Tril(), input_info)\n",
    "    verify_model(InplaceTril(), input_info)\n",
    "\n",
    "\n",
    "def test_triu():\n",
    "    \"\"\"test relay to relax for triu\"\"\"\n",
    "\n",
    "    class Triu(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.triu(data, 1)\n",
    "\n",
    "    class InplaceTriu(Module):\n",
    "        def forward(self, data):\n",
    "            data.triu_(1)\n",
    "            return data\n",
    "\n",
    "    input_info = [([10, 10], \"float32\")]\n",
    "    verify_model(Triu(), input_info)\n",
    "    verify_model(InplaceTriu(), input_info)\n",
    "\n",
    "\n",
    "def test_new_ones():\n",
    "    \"\"\"test relay to relax for new_ones\"\"\"\n",
    "\n",
    "    class NewOnes(Module):\n",
    "        def forward(self, x):\n",
    "            return x.new_ones(1, 2, 3)\n",
    "\n",
    "    input_info = [([1, 2, 3], \"float32\")]\n",
    "    verify_model(NewOnes(), input_info, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_expand():\n",
    "    \"\"\"test relay to relax for expand\"\"\"\n",
    "\n",
    "    class Expand(Module):\n",
    "        def forward(self, x):\n",
    "            return x.expand(4, 2, 3, 4)\n",
    "\n",
    "    input_info = [([1, 2, 3, 4], \"float32\")]\n",
    "    verify_model(Expand(), input_info, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_reduce():\n",
    "    \"\"\"test relay to relax for reduce\"\"\"\n",
    "\n",
    "    # sum\n",
    "    class Sum(Module):\n",
    "        def forward(self, x):\n",
    "            return torch.sum(x, (2, 1))\n",
    "\n",
    "    input_info = [([1, 2, 3, 4], \"float32\")]\n",
    "    verify_model(Sum(), input_info)\n",
    "\n",
    "\n",
    "def test_datatype():\n",
    "    \"\"\"test relay to relax for datatype\"\"\"\n",
    "\n",
    "    input_info = [([1, 2, 3, 4], \"float32\")]\n",
    "\n",
    "    # float\n",
    "    class ToFloat(Module):\n",
    "        def forward(self, x):\n",
    "            return x.float()\n",
    "\n",
    "    verify_model(ToFloat(), input_info, build_target=\"llvm\")\n",
    "\n",
    "    # half\n",
    "    class ToHalf(Module):\n",
    "        def forward(self, x):\n",
    "            return x.half()\n",
    "\n",
    "    verify_model(ToHalf(), input_info)\n",
    "\n",
    "    # type\n",
    "    class Type(Module):\n",
    "        def forward(self, x):\n",
    "            return x.type(torch.float32)\n",
    "\n",
    "    verify_model(Type(), input_info, build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_permute():\n",
    "    \"\"\"test relay to relax for permute\"\"\"\n",
    "\n",
    "    class Permute(Module):\n",
    "        def forward(self, x):\n",
    "            return x.permute(0, 3, 2, 1)\n",
    "\n",
    "    input_info = [([1, 2, 3, 4], \"float32\")]\n",
    "    verify_model(Permute(), input_info)\n",
    "\n",
    "\n",
    "def test_reshape():\n",
    "    \"\"\"test relay to relax for reshape\"\"\"\n",
    "\n",
    "    class Reshape(Module):\n",
    "        def forward(self, x):\n",
    "            return x.reshape(2, 12)\n",
    "\n",
    "    input_info = [([1, 2, 3, 4], \"float32\")]\n",
    "    verify_model(Reshape(), input_info)\n",
    "\n",
    "\n",
    "def test_transpose():\n",
    "    \"\"\"test relay to relax for transpose\"\"\"\n",
    "\n",
    "    class Transpose(Module):\n",
    "        def forward(self, x):\n",
    "            return x.transpose(1, 3)\n",
    "\n",
    "    input_info = [([1, 2, 3, 4], \"float32\")]\n",
    "    verify_model(Transpose(), input_info)\n",
    "\n",
    "\n",
    "def test_view():\n",
    "    \"\"\"test relay to relax for view\"\"\"\n",
    "\n",
    "    class View(Module):\n",
    "        def forward(self, x):\n",
    "            return x.view(2, 12)\n",
    "\n",
    "    input_info = [([1, 2, 3, 4], \"float32\")]\n",
    "    verify_model(View(), input_info)\n",
    "\n",
    "\n",
    "def test_keep_params():\n",
    "    \"\"\"test relay to relax for keep_params\"\"\"\n",
    "\n",
    "    class Conv2D1(Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.conv = torch.nn.Conv2d(3, 6, 7, bias=True)\n",
    "\n",
    "        def forward(self, data):\n",
    "            return self.conv(data)\n",
    "\n",
    "    verify_model(Conv2D1(), [([1, 3, 10, 10], \"float32\")])\n",
    "\n",
    "\n",
    "def test_unwrap_unit_return_tuple():\n",
    "    \"\"\"test relay to relax for unwrap_unit_return_tuple\"\"\"\n",
    "\n",
    "    class Identity(Module):\n",
    "        def forward(self, x):\n",
    "            return (x,)\n",
    "\n",
    "    verify_model(Identity(), [([256, 256], \"float32\")], build_target=\"llvm\")\n",
    "\n",
    "\n",
    "def test_no_bind_return_tuple():\n",
    "    \"\"\"test relay to relax for no_bind_return_tuple\"\"\"\n",
    "\n",
    "    class Identity(Module):\n",
    "        def forward(self, x, y):\n",
    "            return (x, y)\n",
    "\n",
    "    input_info = [([256, 256], \"float32\"), ([256, 256], \"float32\")]\n",
    "    verify_model(Identity(), input_info)\n",
    "\n",
    "\n",
    "def test_argmax():\n",
    "    \"\"\"test relay to relax for argmax\"\"\"\n",
    "\n",
    "    class Argmax1(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.argmax(data, dim=-1)\n",
    "\n",
    "    class Argmax2(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.argmax(data, dim=-1, keepdim=True)\n",
    "\n",
    "    verify_model(Argmax1(), [([256, 256], \"float32\")])\n",
    "    verify_model(Argmax2(), [([256, 256], \"float32\")])\n",
    "\n",
    "\n",
    "def test_to():\n",
    "    \"\"\"test relay to relax for to\"\"\"\n",
    "\n",
    "    class To1(Module):\n",
    "        def forward(self, data):\n",
    "            return data.to(torch.float16)\n",
    "\n",
    "    class To2(Module):\n",
    "        def forward(self, data):\n",
    "            return data.to(\"cpu\")\n",
    "\n",
    "    verify_model(To1(), [([256, 256], \"float32\")])\n",
    "    verify_model(To2(), [([256, 256], \"float32\")])\n",
    "\n",
    "\n",
    "def test_mean():\n",
    "    \"\"\"test relay to relax for mean\"\"\"\n",
    "\n",
    "    class Mean(Module):\n",
    "        def forward(self, data):\n",
    "            return data.mean(-1)\n",
    "\n",
    "    class MeanKeepDim(Module):\n",
    "        def forward(self, data):\n",
    "            return data.mean(-1, keepdim=True)\n",
    "\n",
    "    verify_model(Mean(), [([256, 256], \"float32\")])\n",
    "    verify_model(MeanKeepDim(), [([256, 256], \"float32\")])\n",
    "\n",
    "\n",
    "def test_rsqrt():\n",
    "    \"\"\"test relay to relax for rsqrt\"\"\"\n",
    "\n",
    "    class Rsqrt(Module):\n",
    "        def forward(self, data):\n",
    "            return torch.rsqrt(data)\n",
    "\n",
    "    verify_model(Rsqrt(), [([256, 256], \"float32\")])\n",
    "\n",
    "\n",
    "def test_neg():\n",
    "    \"\"\"test relay to relax for neg\"\"\"\n",
    "\n",
    "    class Neg(Module):\n",
    "        def forward(self, data):\n",
    "            return -data\n",
    "\n",
    "    verify_model(Neg(), [([256, 256], \"float32\")])\n",
    "\n",
    "\n",
    "def test_max():\n",
    "    \"\"\"test relay to relax for max\"\"\"\n",
    "\n",
    "    class Max(Module):\n",
    "        def forward(self, x, y):\n",
    "            return torch.max(x, y)\n",
    "\n",
    "    verify_model(Max(), [([256, 256], \"float32\"), ([256, 256], \"float32\")])\n",
    "\n",
    "\n",
    "def test_name_string_with_colon():\n",
    "    \"\"\"test name string with colons,\n",
    "    e.g., TFLite default input name 'serving_default_input:0'\n",
    "    \"\"\"\n",
    "\n",
    "    dtype = \"float32\"\n",
    "    x_var = relay.var(\"input_0:0\", shape=(3, 5), dtype=dtype)\n",
    "    y_var = relay.var(\"input_1:0\", shape=(3, 5), dtype=dtype)\n",
    "    z_add = relay.add(x_var, y_var)\n",
    "    func = relay.Function([x_var, y_var], z_add)\n",
    "    mod = IRModule()\n",
    "    mod[\"main\"] = func\n",
    "\n",
    "    try:\n",
    "        graph, _ = translate.from_relay(mod)\n",
    "    except Exception as err:\n",
    "        raise RuntimeError(f\"Translation from relay to graph failed: {err}\")\n",
    "    inspect = graph.inspect()\n",
    "\n",
    "    expected = {\n",
    "        \"inputs\": [\n",
    "            {\"name\": \"input_0:0\", \"shape\": [3, 5], \"dtype\": dtype, \"layout\": \"\"},\n",
    "            {\"name\": \"input_1:0\", \"shape\": [3, 5], \"dtype\": dtype, \"layout\": \"\"},\n",
    "        ],\n",
    "        \"outputs\": [{\"name\": \"add\", \"shape\": [3, 5], \"dtype\": dtype, \"layout\": \"\"}],\n",
    "        \"nodes\": {\"total\": 3, \"input\": 2, \"add\": 1},\n",
    "    }\n",
    "\n",
    "    assert msc_utils.dict_equal(inspect, expected), \"Inspect {} mismatch with expected {}\".format(\n",
    "        inspect, expected\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xin",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
